package org.dino.chat.controller;

import org.dino.chat.entity.ActorFilms;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.converter.BeanOutputConverter;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import static java.util.stream.Collectors.joining;

/**
 * Chat Client API
 *
 * @author dino
 * @since 2024/10/28 20:36
 */
@RestController
@RequestMapping("/chat/ai")
public class ChatController {

    private final ChatClient chatClient;

    // 历史消息列表
    static List<Message> historyMessage = new ArrayList<>();
    // 历史消息列表的最大长度
    static int maxLen = 10;


    @Autowired
    public ChatController(ChatClient chatClient) {
        this.chatClient = chatClient;
    }

    @GetMapping("/simple")
    public Map<String, String> completion(@RequestParam(value = "message", defaultValue = "给我讲个笑话") String message) {
        return Map.of("completion", this.chatClient.prompt().user(message).call().content());
    }

    @GetMapping("/simple2")
    Map<String, String> completion(@RequestParam(value = "message", defaultValue = "给我讲个笑话") String message, String voice) {
        return Map.of("completion",
                this.chatClient.prompt()
                        .system(sp -> sp.param("voice", voice))
                        .user(message)
                        .call()
                        .content());
    }


    @GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> stream(@RequestParam(value = "message", defaultValue = "给我讲个笑话") String message, String voice) {
        return this.chatClient.prompt(message)
                .system(sp -> sp.param("voice", voice))
                .stream()
                .content();
    }

    @GetMapping("/actor")
    public ActorFilms actor(@RequestParam(value = "message", defaultValue = "Generate the filmography for a random actor") String message) {
        return this.chatClient.prompt().user(message).call().entity(ActorFilms.class);
    }

    @GetMapping("/actor2")
    public List<ActorFilms> actor2(@RequestParam(value = "message", defaultValue = "Generate the filmography of 5 movies for Tom Hanks and Bill Murray.") String message) {
        return this.chatClient.prompt().user(message).call().entity(new ParameterizedTypeReference<>() {
        });
    }

    @GetMapping("/actor3")
    public List<ActorFilms> actor3(@RequestParam(value = "message", defaultValue = "Generate the filmography for a random actor.") String message) {
        var converter = new BeanOutputConverter<>(new ParameterizedTypeReference<List<ActorFilms>>() {
        });
        Flux<String> flux = this.chatClient.prompt()
                .user(u -> u.text("{message}{format}")
                        .param("format", converter.getFormat())
                        .param("message", message))
                .stream()
                .content();

        String content = flux.collectList().block().stream().collect(joining());

        List<ActorFilms> actorFilms = converter.convert(content);
        return actorFilms;
    }

    @GetMapping("/context")
    public String context(String prompt) {
        // 用户输入的文本是UserMessage
        historyMessage.add(new UserMessage(prompt));
        // 发给AI前对历史消息对列的长度进行检查
        if (historyMessage.size() > maxLen) {
            historyMessage = historyMessage.subList(historyMessage.size() - maxLen - 1, historyMessage.size());
        }
        // 获取AssistantMessage
        ChatResponse chatResponse = this.chatClient.prompt(new Prompt(historyMessage)).call().chatResponse();
        AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
        // 将AI回复的消息放到历史消息列表中
        historyMessage.add(assistantMessage);
        return assistantMessage.getText();
    }

}
